#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA
#

from DynamicSolverBase import DynamicSolverBase

import time
from numpy import c_
from scipy import zeros, asarray
from scipy.sparse.construct import spdiags

from PIDEStatic1D import PIDEStatic1D
from utils.ProjOperators import ProjOnShapeFunctions


class PIDESolver(DynamicSolverBase):

    """
    Integro-partial differential equations solver.
    Finite elements P1 and P2 lagrange.
    """

    __slots__ = ['m_pide', 'm_pival']

    def __init__(self, nsegs, xleft, xright):
        """
        Use the matrices assemble by the static engine to call
        the time stepper.
        """

        DynamicSolverBase.__init__(self)
        self.m_pide = PIDEStatic1D(nsegs, xleft, xright)

        self.m_pival = None

    def __assemble_matrices(self):
        """ 3 Matrices: mass, pde, jump. """

        # mass matrix.

        mass_diag = self.m_pide.get_mass()  # Time step discretization.

        data = [list(mass_diag[1:-1, 0].flatten()), list(mass_diag[1:
                -1, 1].flatten()), list(mass_diag[1:-1, 2].flatten())]
        diags = [0, 1, -1]

        if 3 == self.get_sfunc().dof():
            data.append(list(mass_diag[1:-1, 3].flatten()))
            data.append(list(mass_diag[1:-1, 4].flatten()))
            diags.append(2)
            diags.append(-2)

        mat_size = mass_diag[1:-1, 0].size
        mass_diag = spdiags(asarray(data), asarray(diags), mat_size,
                            mat_size).tocsr()

        # pde matrix.

        pde_diag = self.m_pide.get_pde()  # PDE contribution

        data = [list(pde_diag[1:-1, 0].flatten()), list(pde_diag[1:-1,
                1].flatten()), list(pde_diag[1:-1, 2].flatten())]
        diags = [0, 1, -1]

        if 3 == self.get_sfunc().dof():
            data.append(list(pde_diag[1:-1, 3].flatten()))
            data.append(list(pde_diag[1:-1, 4].flatten()))
            diags.append(2)
            diags.append(-2)

        pde_diag = spdiags(asarray(data), asarray(diags), mat_size,
                           mat_size).tocsr()

        # Jump matrix and boundary operator.

        (jump_full, bop) = self.m_pide.get_jump()

        return (mass_diag, pde_diag, jump_full)

    def __assemble_boperators(self):
        """ Boundary operators are vectors. """

        (jump_full, pide_bop) = self.m_pide.get_jump()

        rhs_size = 1 + pide_bop.size / 2

        # Compute jump contribution.

        jlbop = self.m_pide.get_lambda() * pide_bop[:rhs_size]
        jrbop = self.m_pide.get_lambda() * pide_bop[-rhs_size:]

        jlbop = jlbop[::-1]
        jrbop = jrbop[::-1]

        # Compute Stiff contribution.

        slbop = zeros(jlbop.shape)
        srbop = zeros(jrbop.shape)

        pde_diag = self.m_pide.get_pde()
        if 3 == self.get_sfunc().dof():
            slbop[2] = pde_diag[0, 4]
            srbop[-3] = pde_diag[-1, 3]
        slbop[1] = pde_diag[0, 2]
        srbop[-2] = pde_diag[-1, 1]

        # Mass operator

        mlbop = zeros(slbop.shape, 'float64')
        mrbop = zeros(srbop.shape, 'float64')

        mass_diag = self.m_pide.get_mass()
        if 3 == self.get_sfunc().dof():
            mlbop[2] = mass_diag[0, 4]
            mrbop[-3] = mass_diag[-1, 3]
        mlbop[1] = mass_diag[0, 2]
        mrbop[-2] = mass_diag[-1, 1]

        return (c_[mlbop[1:-1], mrbop[1:-1]], c_[slbop[1:-1], srbop[1:-1]],
                c_[jlbop[1:-1], jrbop[1:-1]])

    def __ready2simulate(self):
        """ Is all what is needed being set. """

        if self.get_ivalf() is None:
            raise ValueError, 'set initial value as a fucntion'

        if self.get_load() is None:
            raise ValueError, 'set load as a function'

    def __proj_ival(self):
        """ Project the initial value in the finite element space. """

        stime = self.get_tstepper().get_startat()

        self.m_pival = ProjOnShapeFunctions(self.m_pide.get_sfunc(),
                self.m_pide.get_quads(), self.m_pide.get_mesh(),
                self.get_ivalf(), stime)

    def _time_step(self):
        """ Set up time stepper, time step and return final result. """

        # Initialize Time stepper.

        tstepper = self.get_tstepper()

        # boundary, ival, load

        (xleft, xright) = (self.m_pide.get_mesh().get_grid()[0],
                           self.m_pide.get_mesh().get_grid()[-1])
        tstepper.set_xlxr(xleft, xright)
        tstepper.set_ival(self.m_pival[1:-1])
        tstepper.set_bcond(self.get_lbc(), self.get_rbc())
        tstepper.set_load(self.get_load())

        # Operators (vector and matrices)

        (mbop, sbop, jbop) = self.__assemble_boperators()
        tstepper.set_boperators(mbop, sbop, jbop)
        (mass, pde, jump) = self.__assemble_matrices()
        tstepper.set_moperators(mass, pde, jump)
        tstepper.set_finite_elem_space(self.m_pide.get_sfunc(),
                self.m_pide.get_quads(), self.m_pide.get_mesh())

        # if 2 == self.m_pide.get_sfunc().dof():
        #    tstepper.set_mat_attrib(3)
        # elif 3 == self.m_pide.get_sfunc().dof():
        #    tstepper.set_mat_attrib(5)
        # else:
        #    raise ValueError, 'only 2 degrees of freedom supported.'

        return tstepper.step()

    def get_mesh(self):
        """ Return the grid used for simulation.  """

        return self.m_pide.get_mesh()

    def get_sfunc(self):
        """ Shape functions defining the finite element space. """

        return self.m_pide.get_sfunc()

    def get_sfunc(self):
        """ Shape functions defining the finite element space. """

        return self.m_pide.get_sfunc()

    def get_pival(self):
        """ Initial value projected in the finite element space.  """

        return self.m_pival

    def set_kernel(self, kernel):
        """ Kernel defining the integral part of the equation. """

        self.m_pide.set_kernel(kernel)

    def set_pide_cos(self, diff, conv, reac, jump):
        """ Diffusion convection reaction jump """

        self.m_pide.set_pide_cos(diff, conv, reac, jump)

    def set_sfunc_and_quads(self, shfunts, quad):
        """ Shape functions and quadrature formulas. """

        self.m_pide.set_sfunc_and_quads(shfunts, quad)

    def solve(self):
        """ Call time stepping engine and post processes resutls. """

        self.__ready2simulate()
        self.__proj_ival()

        return self._time_step()


